# -*- coding:utf-8 -*-
from typing import List

from bson import ObjectId
from loguru import logger

from exceptions import InvalidParam
from exceptions import ModelNoExists
from ioservice.common.mongo import fake_delete_models_by_ids, bulk_update_models_unset_fields_though_id, \
    bulk_update_models_through_id, find_many_docs, find_one_doc, insert_one_model, update_one_model, \
    count_documents as mongo_count_docs, insert_models, check_models_exists_by_ids, aggregate_docs
from ioservice.const import institution_col, institution_counter, activity_col
from ioservice.objectcode import bulk_generate_object_code, generate_object_code
from models.institution import Institution, BaseSearchModel, InstitutionSearchModel, InstitutionUpdateModel
from utils.dateutil import date_now
from utils.iterutil import generate_batch
from utils.regionutil import area2province


def parse_institution_search_query(context):
    exclude_fields = BaseSearchModel.__fields__.keys()
    query = context.dict(exclude=exclude_fields, exclude_unset=True)
    if context.name:
        query.update({"institution_name": {"$regex": f".*{context.name}.*", "$options": "$i"}})
        del query['name']
    if query.get("province"):
        query['region.province'] = context.province
        del query['province']
    if context.query:
        query.update(context.query)
    if context.region:
        if context.region not in area2province:
            raise InvalidParam("region", "not correct region")
        if not query.get("$and"):
            query["$and"] = []
        query["$and"].append({"region.province": {"$in": list(area2province[context.region])}})
        del query["region"]
    if not context.is_deleted:
        query['is_deleted'] = {"$ne": True}
    return query


async def get_institution_served_people_num(institution_ids):
    agg_query = [
        {"$match": {"institution_id": {"$in": institution_ids}, "is_deleted": {"$ne": True}}},
        {
            "$project": {
                "_id": 1,
                "institution_id": 1
            }
        },
        {
            "$lookup": {
                "from": "activity_register",
                "let": {"act_id": "$_id"},
                "pipeline": [
                    {"$match": {"is_deleted": {"$ne": True}}},
                    {"$match": {"$expr": {"$eq": ['$activity_id', "$$act_id"]}}},
                    {"$project": {"person_id": 1}}
                ],
                "as": "registered",
            }
        },
        {
            "$lookup": {
                "from": "alumni",
                "let": {"registerRefId": "$registered"},
                "pipeline": [
                    {"$match": {"$expr": {"$in": ['$_id', "$$registerRefId.person_id"]}}},
                    {"$project": {"_id": 1}}
                ],
                "as": "alumni",
            }
        },
        {
            "$unwind": "$alumni"
        },
        {
            "$group": {
                "_id": "$institution_id",
                "served_people_num": {"$addToSet": "$alumni"}
            }
        },
        {
            "$project": {
                "item": 1,
                "served_people_num": {
                    "$cond": {"if": {"$isArray": "$served_people_num"}, "then": {"$size": "$served_people_num"},
                              "else": "NA"}
                }
            }
        }
    ]
    agg_results = await aggregate_docs(activity_col, agg_query)
    return {str(agg_res["_id"]): agg_res["served_people_num"] for agg_res in agg_results}


async def sync_institutions_data(institutions=None):
    if not institutions:
        institution_docs = await find_many_docs(institution_col, {}, projection=["_id"])
        institutions = [institution['_id'] for institution in institution_docs]
        logger.info(f"institutions not defined: start to update every institutions(total: {len(institutions)})")
    await update_institution_served_people_num(institutions)


async def update_institution_served_people_num(institutions):
    # prepare data
    institutions = [ObjectId(ins_id) for ins_id in institutions]
    served_people_nums = await get_institution_served_people_num(institutions)
    # build update models
    update_models = []
    for ins_id in institutions:
        update_ins = Institution(id=ins_id)
        update_ins.served_people_num = served_people_nums.get(str(ins_id), 0)
        update_models.append(update_ins)
    if not update_models:
        return
    await bulk_update_models_through_id(institution_col, update_models)


async def search_institutions(context: InstitutionSearchModel):
    query = parse_institution_search_query(context)

    query['level'] = {"$ne": 0}
    total = await mongo_count_docs(institution_col, query)
    if not total:
        return {"total": 0, "data": []}
    projection = list(
        set(context.fields) | set(s[0] for s in context.sort)) if context.fields and context.sort else context.fields
    projection = {key: True for key in projection} if projection else {}

    agg_param = [
        {
            "$match": query
        }
    ]
    if context.sort:
        new_sort = list(context.sort)
        add_fields = {}
        for field, direction in context.sort:
            null_field = {f"isNull_{field}": {"$toBool": {"$ifNull": [f'${field}', 0]}}}
            add_fields.update(null_field)
            new_sort = [(f"isNull_{field}", -1)] + new_sort
        sort_params = [
            {"$addFields": add_fields},
            {"$sort": {key: value for key, value in new_sort}},
        ]
        agg_param.extend(sort_params)
        [projection.update({field: False}) for field in add_fields.keys()]
    if context.offset:
        agg_param.append({"$skip": context.offset})
    if context.size:
        agg_param.append({"$limit": context.size})
    else:
        agg_param.append({"$limit": 5000})

    if projection:
        agg_param.append({"$project": projection})

    print(agg_param)
    institutions = await aggregate_docs(institution_col, agg_param)
    for institution in institutions:
        institution["id"] = str(institution.pop("_id"))
    return {"total": total, "data": institutions}


def region_aggregation(province_field="region.province"):
    return {"region": [
        {
            "$project": {
                "province": {
                    "$switch": {
                        "branches": [
                            {
                                "case": {
                                    "$in": [
                                        f"${province_field}",
                                        [
                                            "北京市",
                                            "天津市",
                                            "河北省",
                                            "山西省",
                                            "内蒙古自治区"
                                        ]
                                    ]
                                },
                                "then": "华北地区"
                            },
                            {
                                "case": {
                                    "$in": [
                                        f"${province_field}",
                                        [
                                            "黑龙江省",
                                            "吉林省",
                                            "辽宁省"
                                        ]
                                    ]
                                },
                                "then": "东北地区"
                            },
                            {
                                "case": {
                                    "$in": [
                                        f"${province_field}",
                                        [
                                            "上海市",
                                            "江苏省",
                                            "浙江省",
                                            "安徽省",
                                            "江西省",
                                            "山东省",
                                            "福建省",
                                            "台湾省"
                                        ]
                                    ]
                                },
                                "then": "华东地区"
                            },
                            {
                                "case": {
                                    "$in": [
                                        f"${province_field}",
                                        [
                                            "河南省",
                                            "湖北省",
                                            "湖南省"
                                        ]
                                    ]
                                },
                                "then": "华中地区"
                            },
                            {
                                "case": {
                                    "$in": [
                                        f"${province_field}",
                                        [
                                            "广东省",
                                            "广西壮族自治区",
                                            "海南省",
                                            "香港特别行政区",
                                            "澳门特别行政区"
                                        ]
                                    ]
                                },
                                "then": "华南地区"
                            },
                            {
                                "case": {
                                    "$in": [
                                        f"${province_field}",
                                        [
                                            "重庆市",
                                            "四川省",
                                            "贵州省",
                                            "云南省",
                                            "西藏自治区"
                                        ]
                                    ]
                                },
                                "then": "西南地区"
                            },
                            {
                                "case": {
                                    "$in": [
                                        f"${province_field}",
                                        [
                                            "陕西省",
                                            "甘肃省",
                                            "青海省",
                                            "宁夏回族自治区",
                                            "新疆维吾尔自治区"
                                        ]
                                    ]
                                },
                                "then": "西北地区"
                            }
                        ],
                        "default": "None"
                    }
                }
            }
        },
        {
            "$group": {
                "_id": "$province",
                "num": {
                    "$count": {}
                }
            }
        }
    ]}


def in_kexie_aggregation():
    return {
        "in_kexie": [
            {
                "$group": {
                    "_id": "$in_kexie",
                    "num": {
                        "$sum": 1
                    }
                }
            }
        ]
    }


def get_multi_aggregation_query(query, aggregation_fields: List[str]):
    aggregation_query = [{"$match": query}]
    if not aggregation_fields:
        return {}
    facet = {}
    for field in aggregation_fields:
        if field == "institution_region":
            facet.update(region_aggregation())
        elif field == "in_kexie":
            facet.update(in_kexie_aggregation())
    aggregation_query.append({"$facet": facet})
    return aggregation_query


async def aggregate_institution(params: InstitutionSearchModel):
    query = parse_institution_search_query(params)
    aggregation_query = get_multi_aggregation_query(query, params.aggs or ["institution_region", "in_kexie"])
    result = await aggregate_docs(institution_col, aggregation_query)
    result = [{key: value} for key, value in result[0].items()]
    return result


async def get_institution(institution_name=None, province=None, **kwargs):
    query = {"is_deleted": {"$ne": True}}
    if institution_name:
        query['institution_name'] = institution_name
    if province:
        query['region.province'] = province
    count = await mongo_count_docs(institution_col, query)
    institutions = await find_many_docs(institution_col, query, skip=kwargs.get("skip", 0),
                                        limit=kwargs.get("limit", 20))
    for institution in institutions:
        institution['id'] = str(institution['_id'])
        del institution['_id']
    return {"total": count, "data": institutions}


async def get_institution_by_id(_id, **kwargs):
    institution = await find_one_doc(institution_col, {"_id": ObjectId(_id), "is_deleted": {"$ne": True}}, **kwargs)
    if not institution:
        raise ModelNoExists("institution", str(_id))
    institution["id"] = str(institution["_id"])

    # 查询机构对应的赋能组织数量

    del institution["_id"]
    return institution


async def get_institution_by_query(query, **kwargs):
    institutions = await find_many_docs(institution_col, query, **kwargs)
    for institution in institutions:
        institution["id"] = str(institution.pop("_id"))
    return institutions


async def check_institution_exist_by_ids(ids):
    return await check_models_exists_by_ids(institution_col, ids)


async def insert_institution(institution: Institution):
    # TODO 添加检查
    institution.set_ocode(await generate_object_code(institution_counter))
    institution.create_time = date_now()
    institution.update_time = institution.create_time
    institution_id = await insert_one_model(institution_col, institution)
    return str(institution_id)


async def insert_institutions(institutions: List[Institution], batch_size=100):
    results = set()
    ocodes = await bulk_generate_object_code(institution_counter, len(institutions))
    for ocode, institution in zip(ocodes, institutions):
        institution.set_ocode(ocode)
    for batch in generate_batch(institutions, batch_size=batch_size):
        inserted = await insert_models(institution_col, batch)
        if inserted:
            results.update([str(ins) for ins in inserted])
    return results


async def drop_institution_by_id(_id):
    res = await find_one_doc(institution_col, {"_id": ObjectId(_id), "is_deleted": {"$ne": True}}, projection=["_id"])
    if not res:
        # 组织不存在，抛出异常
        raise ModelNoExists("institution", str(_id))
    res = await fake_delete_models_by_ids(institution_col, [_id])
    return res


async def update_institution(institution_id, institution: Institution):
    # TODO 新增校验
    if not institution_id:
        raise ModelNoExists("institution", "id")

    res = await find_one_doc(institution_col, {"_id": institution_id}, projection=["_id"])
    if not res:
        # 组织不存在，抛出异常
        raise ModelNoExists("institution", str(institution_id))
    del institution.id
    institution.update_time = date_now()
    res = await update_one_model(institution_col, {"_id": institution_id}, institution)
    return res


async def bulk_update_institutions(institutions: List[InstitutionUpdateModel]):
    if len(institutions) == 1:
        return await bulk_update_models_through_id(institution_col, institutions)
    res = await bulk_update_models_unset_fields_though_id(institution_col, institutions, Institution)
    return res


async def crud_institutions_tree(by: bool):
    """获得科协内部机构树"""
    pipe = [
        {"$project": {"in_kexie": 1, "institution_name": 1, "level": 1, "parent_institution_name": 1, "create_time": 1,
                      "update_time": 1,
                      "served_people_num": 1, "hierarchy": 1}},
        {"$facet": {
            "root": [
                {
                    "$match": {
                        "level": 0,
                    }
                }
            ],
            "children": [
                {"$match": {"level": 1, "parent_institution_name": {"$in": ["直属单位", "机关部门"]}}},
                {"$graphLookup": {
                    "from": "institution",
                    "startWith": "$institution_name",
                    "connectFromField": "institution_name",
                    "connectToField": "parent_institution_name",
                    "maxDepth": 0,
                    "as": "hierarchy",
                }},
                {"$sort": {"_id": 1}}
            ]

        },
        },
        {"$unwind": "$root"},
        {"$addFields": {"id": "$_id"}},
        {"$project": {"root._id": 1,
                      "root.institution_name": 1,
                      "root.level": 1,
                      "root.hierarchy": {
                          "$filter": {
                              "input": "$children",
                              "as": "sub_level",
                              "cond": {
                                  "$eq": [
                                      "$$sub_level.parent_institution_name",
                                      "$root.institution_name"
                                  ]
                              }
                          }

                      }}},
        {
            "$replaceRoot": {
                "newRoot": "$root"
            }
        }

    ]
    result = await aggregate_docs(institution_col, pipe)
    for item in result:
        item['id'] = str(item['_id'])
        for h in item['hierarchy']:
            h['id'] = str(h['_id'])
            del h['create_time']
            del h['_id']
            del h['update_time']
            h['hierarchy'] = []
        del item["_id"]
        if by is True:
            item['hierarchy'].append({
                "institution_name": "其他活动数据、其他机构数据、其他人才数据",
                "level": 1,
                "in_kexie": False,
                "parent_institution_name": "其他",
                "hierarchy": [],
                "id": "-1"
            })
    return result


async def crud_institutions_dept(parent_name: str):
    res = await find_many_docs(institution_col, {"parent_institution_name": parent_name, "level": 2},
                               projection=["_id", "level", "parent_institution_name", "institution_name"])
    for r in res:
        r['id'] = str(r['_id'])
        del r['_id']

    return res



async def check_institution_exists_by_ids(ids):
    return await check_models_exists_by_ids(institution_col, ids)
